package org.sean.framework.canary.loadbalancer.selector;

import com.google.common.util.concurrent.ThreadFactoryBuilder;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.client.ServiceInstance;

import java.security.SecureRandom;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 基于响应时间
 */
@Slf4j
public class BestAvailableWeightedSelector extends RoundRobinSelector {
    private final SecureRandom random = new SecureRandom();
    protected ScheduledExecutorService executorService = null;
    protected AtomicBoolean serverWeightAssignmentInProgress = new AtomicBoolean(false);
    ThreadFactory weightTaskThreadFactory = new ThreadFactoryBuilder().setNameFormat("weightTask-pool-%d").build();
    String name = "unknown";
    private int serverWeightTaskInterval = 10;
    // holds the accumulated weight from index 0 to current index
    // for example, element at index 2 holds the sum of weight of servers from 0 to 2
    private volatile Map<String, Double> accumulatedWeights = new HashMap<>();

    public BestAvailableWeightedSelector() {
        super();
    }


    void initialize() {
        if (executorService != null) {
            executorService.shutdown();
        }
        executorService = Executors.newScheduledThreadPool(1, weightTaskThreadFactory);
        executorService.scheduleAtFixedRate(new DynamicServerWeightTask(), 0, serverWeightTaskInterval, TimeUnit.SECONDS);
        // do a initial run
        ServerWeight sw = new ServerWeight();
        sw.maintainWeights();

        Runtime.getRuntime().addShutdownHook(new Thread(() -> executorService.shutdown()));
    }

    public void shutdown() {
        if (executorService != null) {
            executorService.shutdown();
        }
    }

    @Override
    public ServiceInstance choose(String lastServerId, List<ServiceInstance> allServers) {
        ServiceInstance server = null;
        while (server == null) {
            // get hold of the current reference in case it is changed from the other thread
            if (Thread.interrupted()) {
                return null;
            }

            int serverCount = allServers.size();

            if (serverCount == 0) {
                return null;
            }
            try {

                List<Double> currentWeights = new ArrayList<>();
                allServers.forEach(item -> {
                    Double weight = accumulatedWeights.get(item.getUri().toString());
                    if (weight == null || weight < 0.001D) {
                        weight = 1D;
                    }
                    currentWeights.add(weight);
                });
                double allWeight = currentWeights.stream().mapToDouble(Double::doubleValue).sum();

                int serverIndex = 0;
                // generate a random weight between 0 (inclusive) to maxTotalWeight (exclusive)
                double randomWeight = random.nextDouble() * allWeight;
                // pick the server index based on the randomIndex
                int n = 0;
                for (Double d : currentWeights) {
                    if (randomWeight < d) {
                        serverIndex = n;
                        break;
                    } else {
                        randomWeight -= d;
                        n++;
                    }
                }

                server = allServers.get(serverIndex);
                log.debug("test : " + server.getUri().toString());
                if (allServers.size() > 1) {
                    if (server.getUri().toString().equals(lastServerId)) {
                        log.debug("can not choose last server : " + server.getUri());
                        server = null;
                        continue;
                    }
                    return server;
                } else {
                    return server;
                }
            } catch (Exception e) {
                log.warn(e.getMessage(), e);
                return super.choose(lastServerId, allServers);
            }

        }
        return server;
    }

    void setWeights(Map<String, Double> weights) {
        this.accumulatedWeights = weights;
    }

    class DynamicServerWeightTask implements Runnable {
        @Override
        public void run() {
            ServerWeight serverWeight = new ServerWeight();
            try {
                serverWeight.maintainWeights();
            } catch (Exception e) {
                log.error("Error running DynamicServerWeightTask for {}", name, e);
            }
        }
    }

    class ServerWeight {
        public void maintainWeights() {
            if (!serverWeightAssignmentInProgress.compareAndSet(false, true)) {
                return;
            }

//            try {
//                double activeRequestsCount = 0;
//                // 活跃数多代表处理能力差,响应时间长,应该少分配
//                for (ServiceInstance server : getAllServers()) {
//                    // this will automatically load the stats if not in cache
//                    ServerStats ss = stats.getSingleServerStat(server);
//                    // 活跃数
//                    activeRequestsCount += ss.getActiveRequestsCount();
//                }
//                // weight for each server is (sum of responseTime of all servers - responseTime)
//                // so that the longer the response time, the less the weight and the less likely to be chosen
//                Double weightSoFar = 0.0;
//
//                // create new list and hot swap the reference
//                Map<String, Double> finalWeights = new HashMap<>();
//                for (Server server : nlb.getAllServers()) {
//                    ServerStats ss = stats.getSingleServerStat(server);
//                    double weight = activeRequestsCount - ss.getActiveRequestsCount();
//
//                    weightSoFar += weight;
//                    finalWeights.put(server.getHostPort(), weightSoFar);
//                }
//                setWeights(finalWeights);
//            } catch (Exception e) {
//                log.error("Error calculating server weights", e);
//            } finally {
//                serverWeightAssignmentInProgress.set(false);
//            }

        }
    }

}
